﻿#pragma once

#include <learn/alg/alg_pia_curve.h>

namespace xi
{

namespace alg
{

template <std::size_t N> struct Alg_Surface_PIA : Alg_Iterator<std::shared_ptr<Geom_BSplineSurface<N>>>
{
    using Vec = VecT<N, double>; // P

    std::size_t uDegree{3}, vDegree{3};
    std::vector<std::vector<Vec>> data;
    std::vector<std::vector<Vec2>> params;

  protected:
    std::shared_ptr<Geom_BSplineSurface<N>> m_surface;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        if (params.empty())
            params = Alg_Parameter<N>::chord(data);

        std::vector<double> uParams(data.size(), 0);
        std::vector<double> vParams(data[0].size(), 0);

        // 平均参数值得到 u,v 方向的参数序列
        for (std::size_t i = 0; i < data.size(); i++)
        {
            for (std::size_t j = 0; j < data[0].size(); j++)
                uParams[i] += data[i][j].x();
            uParams[i] /= data[i].size();
        }

        for (std::size_t j = 0; j < data[0].size(); j++)
        {
            for (std::size_t i = 0; i < data.size(); i++)
                vParams[j] += data[i][j].y();
            vParams[j] /= data.size();
        }

        auto uKnots = Alg_Knot<N>::average(uParams, uDegree);
        auto vKnots = Alg_Knot<N>::average(vParams, vDegree);

        std::vector<Vec> poles(data.size() * data[0].size());
        for (std::size_t i = 0; i < data.size(); i++)
            for (std::size_t j = 0; j < data[0].size(); j++)
                poles[i * data[0].size() + j] = data[i][j];

        m_surface = std::make_shared<Geom_BSplineSurface<N>>(std::move(poles), std::move(uKnots), std::move(vKnots),
                                                             uDegree, vDegree);
    }

    /**
     * @brief 迭代计算下一个曲面
     *
     * @return std::shared_ptr<Geom_BSplineSurface<N>>
     */
    std::shared_ptr<Geom_BSplineSurface<N>> next() override
    {
        std::vector<Vec> nPoles(data.size() * data[0].size());
        for (std::size_t i = 0; i < data.size(); i++)
        {
            for (std::size_t j = 0; j < data[0].size(); j++)
            {
                auto np = m_surface->pole(i, j) + data[i][j] - m_surface->value(params[i][j].x(), params[i][j].y());
                nPoles[i * data[0].size() + j] = np;
            }
        }

        m_surface = std::make_shared<Geom_BSplineSurface<N>>(std::move(nPoles), m_surface->move_u_knots(),
                                                             m_surface->move_v_knots(), uDegree, vDegree);
        return m_surface;
    }
};

template <std::size_t N> struct Alg_Surface_LSPIA : Alg_Iterator<std::shared_ptr<Geom_BSplineSurface<N>>>
{
    using Vec = VecT<N, double>; // P

    std::size_t uDegree{3}, vDegree{3};
    std::size_t uPoleNum, vPoleNum;
    std::vector<Vec> data;
    std::vector<Vec2> params;
    std::shared_ptr<Geom_BSplineSurface<N>> surface;

  protected:
    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        if (params.empty())
            params = Alg_Parameter<N>::projection(data);

        if (surface == nullptr)
        {
            auto uKnots = Alg_Knot<N>::equally(uPoleNum, uDegree);
            auto vKnots = Alg_Knot<N>::equally(vPoleNum, vDegree);

            std::vector<Vec> poles(uPoleNum * vPoleNum);
            surface = std::make_shared<Geom_BSplineSurface<N>>(std::move(poles), std::move(uKnots), std::move(vKnots),
                                                               uDegree, vDegree);
        }
        else
        {
            uDegree = surface->u_degree();
            vDegree = surface->v_degree();
            uPoleNum = surface->u_pole_num();
            vPoleNum = surface->v_pole_num();
        }
    }

    /**
     * @brief 迭代计算下一个曲面
     *
     * @return std::shared_ptr<Geom_BSplineSurface<N>>
     */
    std::shared_ptr<Geom_BSplineSurface<N>> next() override
    {
        auto m = data.size() - 1;
        auto n1 = uPoleNum - 1;
        auto n2 = vPoleNum - 1;

        std::vector<std::vector<double>> basisSum(n1 + 1,
                                                  std::vector<double>(n2 + 1, 0)); // 记录每个控制点对应的那组基函数和
        std::vector<std::vector<Vec>> deltaSum(n1 + 1,
                                               std::vector<Vec>(n2 + 1)); // 记录每个控制点对应的那组偏差向量加权和

        // 对每个数据点计算增量
        auto uKnots = surface->u_knots();
        auto vKnots = surface->v_knots();
        for (std::size_t i = 0; i < m + 1; i++)
        {
            auto delta = data[i] - surface->value(params[i].x(), params[i].y());

            // 计算此参数对应的非零基函数值
            auto uSpan = find_span(uKnots.data(), uKnots.size(), uDegree, params[i].x());
            auto vSpan = find_span(vKnots.data(), vKnots.size(), vDegree, params[i].y());
            auto uBasis = basis_funcs(uKnots.data(), uDegree, uSpan, params[i].x());
            auto vBasis = basis_funcs(vKnots.data(), vDegree, vSpan, params[i].y());

            // 累计参数
            for (std::size_t j = 0; j < uDegree + 1; j++)
            {
                auto row = uSpan - uDegree + j;
                for (std::size_t k = 0; k < vDegree + 1; k++)
                {
                    auto col = vSpan - vDegree + k;
                    basisSum[row][col] += uBasis[j] * vBasis[k];
                    deltaSum[row][col] += delta * uBasis[j] * vBasis[k];
                }
            }
        }

        for (std::size_t i = 0; i < n1 + 1; i++)
        {
            for (std::size_t j = 0; j < n2 + 1; j++)
            {
                auto pole = surface->pole(i, j);
                auto delta = deltaSum[i][j] / basisSum[i][j];
                surface->set_pole(i, j, pole + delta);
            }
        }
        return surface;
    }
};

template <std::size_t N> struct Alg_Surface_FairingPIA : Alg_Surface_LSPIA<N>
{
    using Vec = VecT<N, double>; // P

    std::vector<double> wStretch; // 每个控制点的拉伸权重
    std::vector<double> wBending; // 每个控制点的光顺权重

  protected:
    // 迭代格式 P' = A * P + B
    Eigen::SparseMatrix<double> m_A;
    Eigen::MatrixXd m_B;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        Alg_Surface_LSPIA<N>::init();

        auto B0 = Alg_CoeffMatrix::surface_least_square(this->params, this->surface->u_knots(),
                                                        this->surface->v_knots(), this->uDegree, this->vDegree);
        auto [B1, B2] = Alg_CoeffMatrix::surface_stretch_bending(this->surface->u_knots(), this->surface->v_knots(),
                                                                 this->uDegree, this->vDegree);
        auto [A, B] = Alg_Curve_FairingPIA<N>::coeff_matrices(B0, B1, B2, this->wStretch, this->wBending, this->data);

        m_A = std::move(A);
        m_B = std::move(B);
    }

    /**
     * @brief 迭代计算下一个曲面
     *
     * @return std::shared_ptr<Geom_BSplineSurface<N>>
     */
    std::shared_ptr<Geom_BSplineSurface<N>> next() override
    {
        auto n1 = this->uPoleNum - 1;
        auto n2 = this->vPoleNum - 1;
        auto n = (n1 + 1) * (n2 + 1);

        // 填充控制点矩阵
        Eigen::MatrixXd P(n, N);
        for (std::size_t i = 0; i < n1 + 1; i++)
        {
            for (std::size_t j = 0; j < n2 + 1; j++)
            {
                auto idx = i * (n2 + 1) + j;
                for (std::size_t k = 0; k < N; k++)
                    P(idx, k) = this->surface->pole(i, j)[k];
            }
        }

        // 计算增量
        P = m_A * P + m_B;
        for (std::size_t i = 0; i < n1 + 1; i++)
        {
            for (std::size_t j = 0; j < n2 + 1; j++)
            {
                auto idx = i * (n2 + 1) + j;
                Vec p;
                for (std::size_t k = 0; k < N; k++)
                    p[k] = P(idx, k);
                this->surface->set_pole(i, j, p);
            }
        }

        return this->surface;
    }
};

} // namespace alg

} // namespace xi
